#pragma once

#include "Calculus.h"

#define ITERATION 100

// 标准线性求解器
class LinearSolver
{
private:
    /**
     * @file EquationSolver.h
     * @author xingyifan
     * @date 2022-07-26 13:33
     *
     * @description: 求解上三角系统
     */

    void solveUpper(Matrix &A, Vector &b, Vector &x)
    {
        int n = b.size();
        for (int i = n - 1; i >= 0; i--)
        {
            Real sum = 0;
            for (int j = i + 1; j < n; j++)
            {
                sum = sum + A[i * n + j] * x[j];
            }
            x[i] = (b[i] - sum) / A[i * n + i];
        }
    }

    /**
     * @file EquationSolver.h
     * @author xingyifan
     * @date 2022-07-26 13:34
     *
     * @description: 求解下三角系统
     */

    void solveLower(Matrix &A, Vector &b, Vector &x)
    {
        int n = b.size();
        x[0] = b[0];
        for (int i = 1; i < n; i++)
        {
            Real sum = 0;
            for (int j = 0; j < i; j++)
            {
                sum = sum + A[i * n + j] * x[j];
            }
            x[i] = b[i] - sum;
        }
    }

    /**
     * @file EquationSolver.h
     * @author xingyifan
     * @date 2022-07-26 13:32
     *
     * @description: 标准 LU 分解
     */

    void LU(Matrix &A)
    {
        int n = A.width();
        for (int i = 0; i < n; i++)
        {
            Real key = A[i * n + i];
            // 第 i 列高斯变换
            for (int j = i + 1; j < n; j++)
            {
                A[j * n + i] = A[j * n + i] / key;
                for (int k = i + 1; k < n; k++)
                {
                    A[j * n + k] = A[j * n + k] - A[i * n + k] * A[j * n + i];
                }
            }
        }
    }

    /**
     * @file EquationSolver.h
     * @author xingyifan
     * @date 2022-07-26 13:32
     *
     * @description: 列主元 LU 分解，返回置换向量
     */

    void LU(Matrix &A, int *perm)
    {
        int n = A.width(); // 矩阵宽度
        Real max, tmp;     // 记录最大值
        int index;         // 记录要交换的主元位置

        for (int i = 0; i < n; i++)
        {
            max = 0;
            index = i;

            // 选取列主元
            for (int p = i; p < n; p++)
            {
                tmp = fabs(A[p * n + i]);
                if (tmp > max)
                {
                    max = tmp;
                    index = p;
                }
            }

            // 在这里奇异,说明 A 的秩为 i
            if (max == 0)
            {
                return;
            }

            // 交换主元
            for (int q = 0; q < n; q++)
            {
                tmp = A[i * n + q];
                A[i * n + q] = A[index * n + q];
                A[index * n + q] = tmp;
            }

            // 记录置换过程
            int tid = perm[i];
            perm[i] = perm[index];
            perm[index] = tid;

            // 正常的高斯消去法
            for (int j = i + 1; j < n; j++)
            {
                A[j * n + i] /= A[i * n + i];
                for (int k = i + 1; k < n; k++)
                {
                    A[j * n + k] -= A[i * n + k] * A[j * n + i];
                }
            }
        }
    }

public:
    /**
     * @file EquationSolver.h
     * @author xingyifan
     * @date 2022-07-26 13:33
     *
     * @description: 求解线性方程组，可选主元 major
     */

    void solve(Matrix &A, Vector &b, bool major = false)
    {
        int n = b.size();

        // 记录置换指标
        std::shared_ptr<int> index(new int[n]);

        if (!major)
        {
            // 标准分解
            LU(A);
        }
        else
        {
            // 初始化置换向量
            for (int i = 0; i < n; i++)
            {
                index.get()[i] = i;
            }

            // 列主元分解
            LU(A, index.get());

            // 置换向量 b
            Vector tb = b;
            for (int i = 0; i < n; i++)
            {
                b[i] = tb[index.get()[i]];
            }
        }

        Vector x(n, 0);
        solveLower(A, b, x);
        solveUpper(A, x, b);
    }
};

/**
 * @file EquationSolver.h
 * @author xingyifan
 * @date 2022-07-26 13:33
 *
 * @description: 求解隐式非线性方程
 */

class FSolve
{
public:
    // 牛顿法求根，可选主元 major
    void solve(FuncX f, Vector &x, bool major = false)
    {
        // Jacobian 算子
        Jacobian j;

        // 线性求解器
        LinearSolver lS;

        // 创建容器
        Vector dx(x.size(), 0);

        for (int i = 0; i < ITERATION; i++)
        {
            // 获得函数值
            Vector fx = f(x);
            dx = fx;

            // 获得 Jacobi 矩阵
            Matrix J = j(f, x);

            // 线性求解
            lS.solve(J, dx, major);
            x = x - dx;

            // 当变化量足够小或者是根就退出
            if (TensorOp::norm(dx) < 1e-8 || TensorOp::norm(fx) < 1e-8)
            {
                break;
            }
        }
    }
};

#undef ITERATION